import nltk
import spacy
import string
import numpy as np
import re
from sklearn.feature_extraction.text import CountVectorizer

nltk.download('punkt')
nltk.download('averaged_perceptron_tagger')

# 初始化 spaCy（只用于句子分割）
nlp = spacy.load('en_core_web_sm', disable=['ner', 'lemmatizer'])


# 1. n-gram频率特征
def extract_ngram_features(texts, n=2, top_k=10):
    """
    提取 top_k 个最常见的 n-gram 特征。
    输入：
        texts: 文本列表（List[str]）
        n: n-gram 的 n 值，例如 2 表示 bi-gram
        top_k: 最多保留的 n-gram 数量
    返回：
        ngram_features: ndarray，每行是一个文本的 n-gram 频率向量
        feature_names: List[str]，对应的 n-gram 特征名称
    """
    vectorizer = CountVectorizer(ngram_range=(n, n), token_pattern=r'\b\w+\b', max_features=top_k)
    ngram_matrix = vectorizer.fit_transform(texts)
    ngram_features = ngram_matrix.toarray()
    feature_names = vectorizer.get_feature_names_out()
    return ngram_features, feature_names


# 2. 标点符号分布特征
def extract_punctuation_features(text):
    """
    分析文本中的标点使用情况。
    返回字典包含总数、密度、常见标点占比等。
    """
    punct_count = sum(text.count(p) for p in string.punctuation)
    text_length = len(text) if len(text) > 0 else 1
    punct_density = punct_count / text_length

    common_puncts = [',', '.', '!', '?', ';']
    punct_ratios = {f'punct_ratio_{p}': text.count(p) / text_length for p in common_puncts}

    features = {
        'punct_count': punct_count,
        'punct_density': punct_density,
    }
    features.update(punct_ratios)
    return features


# 3. 句式复杂度特征
def extract_sentence_features(text):
    """
    统计句子数量、平均长度、从句比例等句式复杂度指标。
    """
    doc = nlp(text)
    sentences = [sent.text for sent in doc.sents]
    sentence_count = len(sentences) if len(sentences) > 0 else 1

    avg_sentence_length = np.mean(
        [len(nltk.word_tokenize(sent)) for sent in sentences]
    ) if sentences else 0

    subordinate_clauses = sum(
        1 for sent in sentences if ',' in sent or 'and' in sent.lower() or 'but' in sent.lower()
    )
    clause_ratio = subordinate_clauses / sentence_count

    return {
        'sentence_count': sentence_count,
        'avg_sentence_length': avg_sentence_length,
        'clause_ratio': clause_ratio
    }
